#include <winsock2.h>
#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#pragma comment(lib, "ws2_32.lib")

void logmsg(const char *msg) {
    FILE *f = fopen("rinetd.log", "a");
    if (f) {
        SYSTEMTIME t;
        GetLocalTime(&t);
        fprintf(f, "[%04d-%02d-%02d %02d:%02d:%02d] %s\n",
                t.wYear, t.wMonth, t.wDay, t.wHour, t.wMinute, t.wSecond,
                msg);
        fclose(f);
    }
}

DWORD WINAPI relay(LPVOID param) {
    SOCKET *socks = (SOCKET *)param;
    char buf[4096];
    int n;
    while ((n = recv(socks[0], buf, sizeof(buf), 0)) > 0) {
        send(socks[1], buf, n, 0);
    }
    closesocket(socks[0]);
    closesocket(socks[1]);
    free(socks);
    return 0;
}

DWORD WINAPI handle_listener(LPVOID param) {
    char **args = (char **)param;
    const char *listen_ip = args[0];
    int listen_port = atoi(args[1]);
    const char *target_ip = args[2];
    int target_port = atoi(args[3]);

    char msg[256];
    SOCKET s = socket(AF_INET, SOCK_STREAM, 0);
    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(listen_port);
    addr.sin_addr.s_addr = inet_addr(listen_ip);

    if (bind(s, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        sprintf(msg, "bind failed on %s:%d", listen_ip, listen_port);
        logmsg(msg);
        return 1;
    }

    listen(s, 5);
    sprintf(msg, "listening on %s:%d", listen_ip, listen_port);
    logmsg(msg);

    while (1) {
        SOCKET client = accept(s, NULL, NULL);
        if (client == INVALID_SOCKET) continue;

        SOCKET remote = socket(AF_INET, SOCK_STREAM, 0);
        struct sockaddr_in target;
        target.sin_family = AF_INET;
        target.sin_port = htons(target_port);
        target.sin_addr.s_addr = inet_addr(target_ip);

        if (connect(remote, (struct sockaddr*)&target, sizeof(target)) < 0) {
            logmsg("connect to remote failed");
            closesocket(client);
            continue;
        }

        logmsg("connection established, starting relay");

        SOCKET *c2r = (SOCKET *)malloc(2 * sizeof(SOCKET));
        SOCKET *r2c = (SOCKET *)malloc(2 * sizeof(SOCKET));
        c2r[0] = client; c2r[1] = remote;
        r2c[0] = remote; r2c[1] = client;

        CreateThread(NULL, 0, relay, c2r, 0, NULL);
        CreateThread(NULL, 0, relay, r2c, 0, NULL);
    }
}

int main(int argc, char *argv[]) {
    WSADATA wsa;
    WSAStartup(MAKEWORD(2,2), &wsa);
    logmsg("rinetd started");

    const char *conf = "rinetd.conf";
    if (argc >= 3 && strcmp(argv[1], "-c") == 0) {
        conf = argv[2];
    }

    FILE *f = fopen(conf, "r");
    if (!f) {
        logmsg("cannot open config file");
        return 1;
    }

    char line[256];
    while (fgets(line, sizeof(line), f)) {
        if (line[0] == '#' || strlen(line) < 7) continue;
        char *listen_ip = strtok(line, " \t\r\n");
        char *listen_port = strtok(NULL, " \t\r\n");
        char *target_ip = strtok(NULL, " \t\r\n");
        char *target_port = strtok(NULL, " \t\r\n");

        if (!listen_ip || !listen_port || !target_ip || !target_port) continue;

        char **args = (char **)malloc(4 * sizeof(char *));
        args[0] = _strdup(listen_ip);
        args[1] = _strdup(listen_port);
        args[2] = _strdup(target_ip);
        args[3] = _strdup(target_port);

        CreateThread(NULL, 0, handle_listener, args, 0, NULL);
    }

    fclose(f);

    while (1) Sleep(10000);
    WSACleanup();
    return 0;
}